[Redshift] DataAPIでトランザクション内での複数クエリ実行ができるようになりました！

全国3兆8000億人のRedshift DataAPIファンの皆様に朗報です！ ついにDataAPIが複数クエリの投入に対応しました！！！

Amazon Redshift Data API now adds support for multi-statement query execution and parameters

え、これまでも複数クエリ投げられたじゃないかって？ 確かに、これまでも複数のSQL分を;で区切って投入することで、複数クエリが動いているようでした。 しかし、これは「そう動いているようだ」というだけであり、 execute-statementのドキュメントにおいて

This statement must be a single SQL statement

execute-statement.html

と記述されているように、複数クエリを投げることは公式には想定されていませんでした。 実際、複数クエリを;で区切って投入した場合、 2つ目以降のクエリは実行はされているように見えますが、結果を取得することが不可能でした。

しかし今回、新しくbatch-execute-statementというAPIが加わり、 複数クエリの投入が正式に可能となりました！ 早速使ってみましょう！

やってみた

AWS CLIからやってみました。 boto3などでも同様のやり方で使えると思います。

AWS CLIコマンドのアップデート

新しく追加されたAPIを使いますので、AWS CLIコマンドをアップデートする必要があります。 私はmacOS上でv2を使っているので、以下のページに沿って行いました。

macOS での AWS CLI バージョン 2 のインストール、更新、アンインストール

$ curl "https://awscli.amazonaws.com/AWSCLIV2.pkg" -o "AWSCLIV2.pkg"
$ sudo installer -pkg AWSCLIV2.pkg -target /

上記コマンドで最新のバージョンにアップデートされました。

$ aws --version
aws-cli/2.2.25 Python/3.8.8 Darwin/20.3.0 exe/x86_64 prompt/off

クエリを投げてみる

複数クエリを使う場合には新しく用意されたbatch-execute-statementを使用します。 指定する引数は従来のexecute-statementと同様ですが、 --sqlに複数のクエリが渡せるようになっています。

ちょっと注意が必要なのは、よくSQLで複数クエリを流すような;で区切る形式ではなく、 SQLを１つずつ引数として渡すという点です。 例を見た方が簡単ですね。

$ aws redshift-data batch-execute-statement \
    --cluster-identifier cm-hirano-redshift-01 \
    --database dev \
    --db-user sample_user \
    --sql "select 1" "select 2"

不定数の引数が渡せるという形なので、--sqlは最後に配置するのが安全そうです （私が観測した範囲では途中に配置しても正しく解釈してくれたようには見えました）。

クエリの情報を取得する

describe-statementを使用して、投入したクエリの情報を取得して見ます。 SubStatementsという項目の中に投入したクエリの数分の情報が入っているのがわかります。

$ aws redshift-data describe-statement --id cd0bb3ba-3d22-41a5-9782-f1d0b764f4b4
{
    "ClusterIdentifier": "cm-hirano-redshift-01",
    "CreatedAt": "2021-08-03T12:08:12.354000+09:00",
    "Duration": 5057362,
    "HasResultSet": true,
    "Id": "cd0bb3ba-3d22-41a5-9782-f1d0b764f4b4",
    "RedshiftPid": 18313,
    "RedshiftQueryId": 0,
    "ResultRows": -1,
    "ResultSize": -1,
    "Status": "FINISHED",
    "SubStatements": [
        {
            "CreatedAt": "2021-08-03T12:08:12.626000+09:00",
            "Duration": 2346589,
            "HasResultSet": true,
            "Id": "cd0bb3ba-3d22-41a5-9782-f1d0b764f4b4:1",
            "QueryString": "select 1",
            "RedshiftQueryId": -1,
            "ResultRows": 1,
            "ResultSize": 11,
            "Status": "FINISHED",
            "UpdatedAt": "2021-08-03T12:08:13.173000+09:00"
        },
        {
            "CreatedAt": "2021-08-03T12:08:12.626000+09:00",
            "Duration": 2710773,
            "HasResultSet": true,
            "Id": "cd0bb3ba-3d22-41a5-9782-f1d0b764f4b4:2",
            "QueryString": "select 2",
            "RedshiftQueryId": -1,
            "ResultRows": 1,
            "ResultSize": 11,
            "Status": "FINISHED",
            "UpdatedAt": "2021-08-03T12:08:13.257000+09:00"
        }
    ],
    "UpdatedAt": "2021-08-03T12:08:13.276000+09:00"
}

各クエリごとにIdが発行されている（親IDの末尾に:数字をつけた形のようです）ので、 このIdでクエリの結果が取得できます。 コマンドは従来と同様get-statement-resultです。

$ aws redshift-data get-statement-result --id 604a771d-08d8-4d4e-8db8-40251919114c:1
{
    "Records": [
        [
            {
                "longValue": 1
            }
        ]
    ],
    "ColumnMetadata": [
        {
            "isCaseSensitive": false,
            "isCurrency": false,
            "isSigned": true,
            "label": "?column?",
            "length": 0,
            "name": "?column?",
            "nullable": 1,
            "precision": 10,
            "scale": 0,
            "schemaName": "",
            "tableName": "",
            "typeName": "int4"
        }
    ],
    "TotalNumRows": 1
}

めでたく、最初のクエリの結果が取得できました！

トランザクションが使える

さて、ここまでの内容だけだと、 「従来のやり方を2回やったのと何も変わらないよね？」と思ってしまいがちです。 しかし、そんなことはないです！

注目すべきは、batch-execute-statementで渡された複数クエリは 単一トランザクションとして実行される という点です！

従来のexecute-statementでは、例え;で区切って複数クエリを投げたとしても、 各クエリはそれぞれ別トランザクションとして実行されているように見えました。 なので、setでトランザクションの設定をした上でのクエリ実行などができませんでした。 しかしbatch-execute-statementを使えばこの問題も解決です！

search_pathにスキーマを追加して、そのスキーマのテーブル一覧を取得するクエリを投げてみます。

$ aws redshift-data batch-execute-statement \
    --cluster-identifier cm-hirano-redshift-01 \
    --database dev \
    --db-user sample_user \
    --sql "SET search_path TO '\$user', 'public', 'sample_schema'" \
          "SELECT pg_table_def.tablename FROM pg_table_def
             WHERE schemaname = 'sample_schema'
             ORDER BY tablename"

SubStatementsのIdは:数字で指定できることがわかっているので、 describe-statementは確認せず、 出力されたIdの最後に:2をつけて2つ目のクエリの結果を取得してみます。

$ aws redshift-data get-statement-result --id 5656243b-59f4-470b-bb36-43ea779e9065:2
{
    "Records": [
        [ { "stringValue": "table_01" } ],
        [ { "stringValue": "table_02" } ],
        （以下省略）
    ],
    "ColumnMetadata": [
        {
            "isCaseSensitive": true,
            "isCurrency": false,
            "isSigned": false,
            "label": "tablename",
            "length": 0,
            "name": "tablename",
            "nullable": 1,
            "precision": 64,
            "scale": 0,
            "schemaName": "pg_catalog",
            "tableName": "pg_table_def",
            "typeName": "name"
        }
    ],
    "TotalNumRows": 112
}

search_pathに追加したsample_schemaスキーマのテーブルも取得できました！！ 従来のDataAPIではこういったクエリは実行できなかったので、 これができるようになって非常に嬉しいです！

コマンドラインからpsqlっぽく実行したい

以前、

[Redshift] DataAPIを使ってpsqlっぽくSQL実行するコマンドを作ってみた。

というブログを書きました。 DataAPIを使ってクエリを実行して、クエリ完了を同期的に待って結果を表示するコマンドです。

こちらのコマンドも今回のアップデートに合わせて、トランザクションでの複数クエリ投入に対応致しました。 具体的には、

$ redshift -C <ClusterName> -D <Database> -U <User> \
    -c 'select 1; select 2;'
1
2

のように;区切りで複数クエリを投入でき、同一トランザクション内で実行されるような形になっております。 クエリはファイルを指定しての実行にも対応しています。 興味があればこちらも活用頂ければと思います！

まとめ

Redshift DataAPIで複数クエリが投げられるbatch-execute-statementがサポートされました。 これで複数クエリを一度に投げることができ、結果も各クエリごとに取得することができます！

また複数クエリはトランザクションの中で実行されるため、 従来のDataAPIでは不可能だったセッションの設定を行った後のクエリ実行が可能になりました！

DataAPIを使うことで、 プライベートサブネットにあるRedshiftへも踏み台を経由することなく （私の場合は、このおかげで会社のVPNに繋ぐ必要がない！）クエリを投げることができるので非常に重宝しています。 しかしこれまではトランザクションがないために一部のクエリは正しく動作させることができず、 泣く泣く踏み台経由で接続したりしていました。 しかしそれももう過去のことです！全てのクエリ実行をDataAPIで行うことも夢ではなくなりました！ ガンガン活用していきましょう！！